/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/
#ifndef itkMultivariateLegendrePolynomial_h
#define itkMultivariateLegendrePolynomial_h

#include "itkIntTypes.h"
#include "itkIndent.h"
#include <vector>
#include "itkArray.h"
#include "ITKPolynomialsExport.h"

#include <sycl/sycl.hpp>
#include <bisheng/bisheng.hpp>

using namespace sycl;
using namespace bisheng;

namespace itk
{
/**
 *\class MultivariateLegendrePolynomial
 * \brief 2D and 3D multivariate Legendre Polynomial
 *
 * In 2D,
 * \f[
 *  f(x_{vector}, parameter_{vector}) =
 * \sum_i^l \left(
 * \sum_j^{l-i} \left( parameter_ {ij} * P_i(x) *P_j(y)) \right) \right)
 * \f]
 * where P_i() denoting a Legendre polynomial of degree i
 * and l it the degree of the polynomial
 *
 * In 3D,
 * \f[
 * f(x_{vector}, parameter_{vector}) =
 * \sum_i^l \left( \sum_j^{l-i} \left( \sum_k^{l-i-j} \left(
 * parameter_{ijk} * P_i(x) * P_j(y) * P_k(z) \right) \right) \right)
 * \f]
 *
 * The size of the parameter vector for 2D is
 * \f$\frac{(l+1)\cdot(1+2)}{2}\f$,
 * and for 3D is \f$\frac{(l+1)*(l+2)*(l+3)}{3!}\f$
 *
 * To get the size of the parameter vector, users can use one of the
 * two GetNumberOfCoefficients() member functions
 *
 * To get function result, users can use the operator() or its
 * SimpleForwardIterator's Get() method.
 *
 * This is a part of the bias correction methods and implementations
 * that was initially developed and implemented
 * by Martin Styner, Univ. of North Carolina at Chapel Hill, and his
 * colleagues.
 *
 * \note For more details. refer to the following articles.
 * "Parametric estimate of intensity inhomogeneities applied to MRI"
 * Martin Styner, G. Gerig, Christian Brechbuehler, Gabor Szekely,
 * IEEE TRANSACTIONS ON MEDICAL IMAGING; 19(3), pp. 153-165, 2000,
 * (http://www.cs.unc.edu/~styner/docs/tmi99.pdf)
 *
 * "Evaluation of 2D/3D bias correction with 1+1ES-optimization"
 * Martin Styner, Prof. Dr. G. Gerig (IKT, BIWI, ETH Zuerich), TR-197
 * (http://www.cs.unc.edu/~styner/docs/StynerTR97.pdf)
 * \ingroup ITKPolynomials
 */

class ITKPolynomials_EXPORT MultivariateLegendrePolynomial
{
public:
  using Self = MultivariateLegendrePolynomial;

  using DoubleArrayType = std::vector<float>;
  using ULongArrayType = std::vector<unsigned long>;
  using LongArrayType = std::vector<long>;

  /** Internal coefficient storage type. */
  using CoefficientArrayType = DoubleArrayType;

  /** Same as CoefficientArray
   * This type definition will be used by EnergyFunction object. */
  using ParametersType = Array<float>;

  /** The size of the domain. */
  using DomainSizeType = ULongArrayType;
  using IndexType = LongArrayType;

  /** Constructor. */
  // MultivariateLegendrePolynomial(unsigned int dimension, unsigned int degree, const DomainSizeType & domainSize);
  MultivariateLegendrePolynomial(unsigned int           dimension,
                                                               unsigned int           degree,
                                                               const DomainSizeType & domainSize)
  {
    if (dimension > 3 || dimension < 2)
    {
      itkGenericExceptionMacro(<< "MultivariateLegendrePolynomial only supports 2D and 3D");
    }
    
    m_Dimension = dimension;
    m_Degree = degree;
    m_DomainSize = domainSize;
    m_NumberOfCoefficients = this->GetNumberOfCoefficients(dimension, degree);

    // used as intermediate store to hold legendre polynomials
    // y_coef[i,j] = Sum (0 <= k <= m-i-j) p(i,j,k) * P(z)
    m_CachedYCoef.resize(((m_Degree + 1) * (m_Degree + 2)) / 2, 0.0);

    // used as intermediate store to hold legendre polynomials
    // x_coef[i] = Sum (0 <= j <= m-i) y_coef[i,j] * P(y)
    m_CachedXCoef.resize(m_Degree + 1, 0.0);

    m_CachedZCoef.resize(m_Degree + 1, 0.0);

    m_CoefficientArray.resize(m_NumberOfCoefficients, 0.0);

    m_PrevY = -1;
    m_PrevZ = -1;

    m_NormFactor = DoubleArrayType(m_Dimension);
    for (unsigned int j = 0; j < m_Dimension; j++)
    {
      m_NormFactor[j] = 2.0f / (static_cast<double>(m_DomainSize[j]) - 1.0f);
    }
  }
  /** Destructor. */
  virtual ~MultivariateLegendrePolynomial()= default;

  /** Gets the dimension. */
  unsigned int
  GetDimension() const
  {
    return m_Dimension;
  }

  /** Gets the degree (the degree of Legendre polynomials). */
  unsigned int
  GetDegree() const
  {
    return m_Degree;
  }

  /** Returns the number of coefficients of the polynomial
   *  This number is computed from the degree of the polynomial
   *  the SetCoefficients() method expects an array of this
   *  size, an exception is thrown otherwise
   *  \sa SetCoefficients
   */
  unsigned int
  GetNumberOfCoefficients() const
  {
    return m_NumberOfCoefficients;
  }

  /** Gets each dimension's size. */
  const DomainSizeType &
  GetDomainSize() const
  {
    return m_DomainSize;
  }

  /**
   *\class CoefficientVectorSizeMismatch
   * \brief Exception object.
   * \ingroup ITKPolynomials
   */
  class CoefficientVectorSizeMismatch
  {
  public:
    CoefficientVectorSizeMismatch(const std::size_t given, const std::size_t required)
      : m_Required{ required }
      , m_Given{ given }
    {}

    std::size_t m_Required;
    std::size_t m_Given;
  };

  /** \brief Sets the Legendre polynomials' parameters.
   * \warning The number of coefficients provided should
   * match the number returned by GetNumberOfCoefficients()
   * otherwise an exception is thrown.  */
  void SetCoefficients(const CoefficientArrayType & coefficients)
  {
    if (coefficients.size() != m_NumberOfCoefficients)
    {
      throw CoefficientVectorSizeMismatch(coefficients.size(), m_NumberOfCoefficients);
    }

    // copy coefficients to array of double
    m_CoefficientArray.resize(m_NumberOfCoefficients);
    for (unsigned int i = 0; i < m_NumberOfCoefficients; i++)
    {
      m_CoefficientArray[i] = coefficients[i];
    }

    m_PrevY = -1;
    m_PrevZ = -1;
  }

  void SetCoefficients(const ParametersType & coefficients)
  {
    if (coefficients.size() != m_NumberOfCoefficients)
    {
      throw CoefficientVectorSizeMismatch(coefficients.size(), m_NumberOfCoefficients);
    }

    // copy coefficients to array of double
    m_CoefficientArray.resize(m_NumberOfCoefficients);
    for (unsigned int i = 0; i < m_NumberOfCoefficients; i++)
    {
      m_CoefficientArray[i] = coefficients[i];
    }

    m_PrevY = -1;
    m_PrevZ = -1;
  }

  /** \brief Gets Legendre polynomials' coefficients. */
  const CoefficientArrayType &
  GetCoefficients() const;

  /** In the case which the bias field is 2D, it returns bias value at
   * the point which is specified by the index */
  double
  Evaluate(IndexType & index)
  {
    if (m_Dimension == 2)
    {
      if (index[1] != m_PrevY)
      {
        // normalized y [-1, 1]
        double norm_y = m_NormFactor[1] * static_cast<double>(index[1] - 1);
        this->CalculateXCoef(norm_y, m_CoefficientArray);
        m_PrevY = index[1];
      }

      // normalized x [-1, 1]
      double norm_x = m_NormFactor[0] * static_cast<double>(index[0] - 1);

      return LegendreSum(norm_x, m_Degree, m_CachedXCoef);
    }
    else if (m_Dimension == 3)
    {
      if (index[2] != m_PrevZ)
      {
        // normalized z [-1, 1]
        double norm_z = m_NormFactor[2] * static_cast<double>(index[2] - 1);
        this->CalculateYCoef(norm_z, m_CoefficientArray);

        m_PrevZ = index[2];
      }
      if (index[1] != m_PrevY)
      {
        // normalized y [-1, 1]
        double norm_y = m_NormFactor[1] * static_cast<double>(index[1] - 1);

        this->CalculateXCoef(norm_y, m_CachedYCoef);
  
        m_PrevY = index[1];
      }
      // normalized x [-1, 1]
      double norm_x = m_NormFactor[0] * static_cast<double>(index[0] - 1);
      return this->LegendreSum(norm_x, m_Degree, m_CachedXCoef);
    }
    return 0;
  }

  /** Gets the number of coefficients. */
  unsigned int
  GetNumberOfCoefficients(unsigned int dimension, unsigned int degree);

  /**
   *\class SimpleForwardIterator
   * \brief Iterator which only supports forward iteration and
   * Begin(), IsAtEnd(), and Get() method which work just like as
   * SimpleImageRegionIterator.
   * \ingroup ITKPolynomials
   */
  class SimpleForwardIterator
  {
  public:
    SimpleForwardIterator(MultivariateLegendrePolynomial * polynomial)
      : m_MultivariateLegendrePolynomial(polynomial)
      , m_Dimension(m_MultivariateLegendrePolynomial->GetDimension())
      , m_DomainSize(m_MultivariateLegendrePolynomial->GetDomainSize())
      , m_IsAtEnd(false)
    {
      m_Index.resize(m_Dimension);
      std::fill(m_Index.begin(), m_Index.end(), 0);
    }

    void
    Begin()
    {
      m_IsAtEnd = false;
      for (unsigned int dim = 0; dim < m_Dimension; dim++)
      {
        m_Index[dim] = 0;
      }
    }

    bool
    IsAtEnd() const
    {
      return m_IsAtEnd;
    }

    SimpleForwardIterator &
    operator++()
    {
      for (unsigned int dim = 0; dim < m_Dimension; dim++)
      {
        if (m_Index[dim] < static_cast<int>(m_DomainSize[dim] - 1))
        {
          m_Index[dim] += 1;
          return *this;
        }
        else
        {
          if (dim == m_Dimension - 1)
          {
            m_IsAtEnd = true;
            break;
          }
          else
          {
            m_Index[dim] = 0;
          }
        }
      }
      return *this;
    }

    double
    Get()
    {
      return m_MultivariateLegendrePolynomial->Evaluate(m_Index);
    }

  private:
    MultivariateLegendrePolynomial * m_MultivariateLegendrePolynomial;
    unsigned int                     m_Dimension;
    DomainSizeType                   m_DomainSize;
    IndexType                        m_Index;
    bool                             m_IsAtEnd;
  }; // end of class Iterator

  void
  Print(std::ostream & os) const
  {
    itk::Indent indent(4);

    this->PrintSelf(os, indent);
  }

protected:
  void
  PrintSelf(std::ostream & os, Indent indent) const
  {
    os << indent << "Dimension: " << m_Dimension << std::endl;
    os << indent << "Degree: " << m_Degree << std::endl;
    os << indent << "DomainSize: ";
    for (const auto i : m_DomainSize)
    {
      os << i << " ";
    }
    os << std::endl;

    os << indent << "Cached X coefficients: ";
    for (auto i : m_CachedXCoef)
    {
      os << i << " ";
    }
    os << std::endl;

    os << indent << "Cached Y coefficients: ";
    for (auto i : m_CachedYCoef)
    {
      os << i << " ";
    }
    os << std::endl;

    os << indent << "Cached Z coefficients: ";
    for (auto i : m_CachedZCoef)
    {
      os << i << " ";
    }
    os << std::endl;

    os << indent << "Coefficients: ";
    for (auto i : m_CoefficientArray)
    {
      os << i << " ";
    }
    os << std::endl;

    os << indent << "Normalization factors: ";
    for (auto i : m_NormFactor)
    {
      os << i << " ";
    }
    os << std::endl;

    os << indent << "Previous Y index: " << m_PrevY << std::endl;
    os << indent << "Previous Z index: " << m_PrevZ << std::endl;
  }

  float LegendreSum(const float x, int n, const CoefficientArrayType & coef, int offset = 0)
  // n+1 elements !
  // nPn(x)=(2n-1)xPn-1(x)-(n-1)Pn-2(x)
  {
    if (n == 0)
    {
      return coef[offset];
    }

    float ykp2 = 0, ykp1 = coef[n + offset];

    for (int k = n - 1; k > 0; k--)
    {
      float yk = x * ykp1 * (2 * k + 1) / (k + 1) - ykp2 * (k + 1) / (k + 2) + coef[k + offset];
      ykp2 = ykp1;
      ykp1 = yk;
    }
    return -ykp2 / 2 + x * ykp1 + coef[offset];
  }

  float LegendreSum(const float x, int n, const float* coef, int offset = 0)
  // n+1 elements !
  // nPn(x)=(2n-1)xPn-1(x)-(n-1)Pn-2(x)
  {
    if (n == 0)
    {
      return coef[offset];
    }

    float ykp2 = 0, ykp1 = coef[n + offset];

    for (int k = n - 1; k > 0; k--)
    {
      float yk = x * ykp1 * (2 * k + 1) / (k + 1) - ykp2 * (k + 1) / (k + 2) + coef[k + offset];
      ykp2 = ykp1;
      ykp1 = yk;
    }
    return -ykp2 / 2 + x * ykp1 + coef[offset];
  }

  void CalculateXCoef(float norm_y, const CoefficientArrayType & coef)
  {
    // compute x_coef[i] = sum (0 <= j <= m-i) pij * P(y)]

    int offset = 0;
    for (unsigned int lx = 0; lx <= m_Degree; lx++)
    {
      m_CachedXCoef[lx] = LegendreSum(norm_y, m_Degree - lx, coef, offset);
      offset += (m_Degree + 1 - lx);
    }
  }

  void CalculateYCoef(double norm_z, const CoefficientArrayType & coefficients_float)
  {
    // compute y_coef[i,j] = sum (0 <= k <= m-i-j) pijk * P(z)
    float norm_zfloat=norm_z;
    queue Q(ascend_selector{});
    int coefSize=coefficients_float.size();
    auto *coefBuf=malloc_device<float> (coefSize,Q);
    Q.memcpy(coefBuf,coefficients_float.data(),coefSize*sizeof(float));
    auto *m_CachedYCoefBuf = malloc_device<float>(m_CachedYCoef.size(), Q);
    int m_CachedZCoef_size=m_CachedZCoef.size();
    int my_Degree=this->m_Degree;
    int B=4;
    int N=this->m_Degree+1;
    int MAX_UB=65536;
    int group_num=N/B==0?1:N/B;

    Q.launch<class CalY>(group_num,[coefSize,MAX_UB,coefBuf,B,N,my_Degree,
                  m_CachedZCoef_size,m_CachedYCoefBuf,norm_zfloat,this](group<1> mygroup){
          int coefSize2= (coefSize%32==0) ? coefSize : 32*(coefSize/32+1);
          int groupId=mygroup.get_group_id();
          int start=groupId*B;
          int end=0;
          if(N/B==0){
            end=start+N;
          }else if( (N%B!=0) && (groupId==N/B-1) ){
            end=start+N-(N/B-1)*B;
          }else{
            end=start+ B;
          }
          int m_Degree=my_Degree;
          __local float m_CachedZCoefBuf_local[m_CachedZCoef_size];

          if(coefSize2<=MAX_UB){
            __local float coefBuf_local[coefSize2];
            sycl::dmi::memcpy_blocks(coefBuf_local,coefBuf,coefSize2*sizeof(float)/32);
            for(int lx = start; lx < end; lx++){
              for(int ly=0;ly<m_Degree+1;ly++){
                if(ly<=m_Degree-lx){
                  unsigned int c_index=lx * (m_Degree * 2 + 3 - lx) / 2+ly;
                  unsigned int y_index=c_index;
                  unsigned int z_index=c_index;
                  unsigned int lzmax=m_Degree-lx-ly;
                  for (unsigned int lz=0;lz<=lzmax;lz++){
                    m_CachedZCoefBuf_local[lz]=coefBuf_local[z_index];
                    z_index+=((m_Degree + 1 - lz) * (m_Degree + 2 - lz) / 2) - lx;
                  }
                  m_CachedYCoefBuf[y_index]=this->LegendreSum(norm_zfloat, m_Degree - lx - ly, m_CachedZCoefBuf_local,0);//20ms
                }
              }
            }
          }else{
            for(int lx = start; lx < end; lx++){
              for(int ly=0;ly<m_Degree+1;ly++){
                if(ly<=m_Degree-lx){
                  unsigned int c_index=lx * (m_Degree * 2 + 3 - lx) / 2+ly;
                  unsigned int y_index=c_index;
                  unsigned int z_index=c_index;
                  unsigned int lzmax=m_Degree-lx-ly;
                  for (unsigned int lz=0;lz<=lzmax;lz++){
                    m_CachedZCoefBuf_local[lz]=coefBuf[z_index];
                    z_index+=((m_Degree + 1 - lz) * (m_Degree + 2 - lz) / 2) - lx;
                  }
                  m_CachedYCoefBuf[y_index]=this->LegendreSum(norm_zfloat, m_Degree - lx - ly, m_CachedZCoefBuf_local,0);//20ms
                }
              }
            }
          }
      });

    Q.memcpy(m_CachedYCoef.data(), m_CachedYCoefBuf, m_CachedYCoef.size() * sizeof(float));
    Q.wait(); 

    free(coefBuf,Q);
    free(m_CachedYCoefBuf,Q);
  }

private:
  DomainSizeType m_DomainSize;
  unsigned int   m_Dimension;
  unsigned int   m_Degree;
  unsigned int   m_NumberOfCoefficients;

  CoefficientArrayType m_CoefficientArray;
  CoefficientArrayType m_CachedXCoef;
  CoefficientArrayType m_CachedYCoef;
  CoefficientArrayType m_CachedZCoef;

  DoubleArrayType m_NormFactor;
  IndexValueType  m_PrevY;
  IndexValueType  m_PrevZ;
}; // end of class

ITKPolynomials_EXPORT std::ostream &
                      operator<<(std::ostream & os, const MultivariateLegendrePolynomial & poly);
} // end of namespace itk
#endif